﻿import pygame
import math
import sys
from pygame.locals import *
import random

# 初始化 Pygame
pygame.init()

# 屏幕设置
WIDTH, HEIGHT = 1200, 800
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("太阳系模拟")

# 颜色定义
BLACK = (0, 0, 0)
YELLOW = (255, 255, 0)
WHITE = (255, 255, 255)
GRAY = (200, 200, 200)
BROWN = (165, 113, 78)
RED = (255, 0, 0)
ORANGE = (255, 165, 0)
BLUE = (0, 0, 255)
LIGHT_BLUE = (173, 216, 230)
DARK_BLUE = (0, 0, 139)
GREEN = (0, 128, 0)

# 时钟控制
clock = pygame.time.Clock()
FPS = 60

# 天文单位 (AU) 到像素的转换比例
AU = 149.6e6 * 1000  # 149.6 million km in meters
SCALE = 200 / AU  # 缩放比例
TIME_STEP = 3600 * 24  # 1 day in seconds

# 行星数据
class CelestialBody:
    def __init__(self, name, radius, color, distance_au, orbital_period, rotation_period, has_rings=False, ring_color=None):
        self.name = name
        self.radius = radius
        self.color = color
        self.distance = distance_au * AU  # 转换为米
        self.orbital_period = orbital_period  # 地球日
        self.rotation_period = rotation_period  # 地球日
        self.angle = 0
        self.rotation_angle = 0
        self.has_rings = has_rings
        self.ring_color = ring_color
        self.orbit_points = []
        self.max_orbit_points = 200
        
    def update(self):
        # 更新公转角度
        self.angle += (2 * math.pi) / (self.orbital_period * TIME_STEP / TIME_STEP)
        
        # 更新自转角度
        self.rotation_angle += (2 * math.pi) / (self.rotation_period * TIME_STEP / TIME_STEP)
        
        # 记录轨道点
        x = self.distance * math.cos(self.angle) * SCALE + WIDTH // 2
        y = self.distance * math.sin(self.angle) * SCALE + HEIGHT // 2
        self.orbit_points.append((x, y))
        
        # 限制轨道点数量
        if len(self.orbit_points) > self.max_orbit_points:
            self.orbit_points.pop(0)
    
    def draw(self, surface):
        # 绘制轨道
        if len(self.orbit_points) > 1:
            pygame.draw.lines(surface, (50, 50, 50), False, self.orbit_points, 1)
        
        # 计算当前位置
        x = self.distance * math.cos(self.angle) * SCALE + WIDTH // 2
        y = self.distance * math.sin(self.angle) * SCALE + HEIGHT // 2
        
        # 绘制行星
        pygame.draw.circle(surface, self.color, (int(x), int(y)), self.radius)
        
        # 绘制行星自转效果
        if self.rotation_period > 0:
            # 绘制一条线表示自转
            end_x = x + self.radius * math.cos(self.rotation_angle)
            end_y = y + self.radius * math.sin(self.rotation_angle)
            pygame.draw.line(surface, (255, 255, 255), (x, y), (end_x, end_y), 1)
        
        # 绘制行星环（如果有）
        if self.has_rings and self.ring_color:
            ring_radius = self.radius * 1.5
            pygame.draw.circle(surface, self.ring_color, (int(x), int(y)), ring_radius, 1)
            pygame.draw.circle(surface, self.ring_color, (int(x), int(y)), ring_radius + 5, 1)
        
        # 绘制行星名称
        font = pygame.font.SysFont(None, 20)
        text = font.render(self.name, True, WHITE)
        surface.blit(text, (x - text.get_width() // 2, y - self.radius - 20))

# 创建太阳和行星
sun = CelestialBody("太阳", 20, YELLOW, 0, 0, 27, False)
mercury = CelestialBody("水星", 5, GRAY, 0.39, 88, 59)
venus = CelestialBody("金星", 8, BROWN, 0.72, 225, 243)
earth = CelestialBody("地球", 9, BLUE, 1.0, 365, 1)
mars = CelestialBody("火星", 7, RED, 1.52, 687, 1.03)
jupiter = CelestialBody("木星", 15, ORANGE, 5.2, 4333, 0.41)
saturn = CelestialBody("土星", 13, LIGHT_BLUE, 9.54, 10759, 0.45, True, (210, 180, 140))
uranus = CelestialBody("天王星", 10, DARK_BLUE, 19.2, 30687, 0.72)
neptune = CelestialBody("海王星", 10, BLUE, 30.06, 60190, 0.67)

planets = [mercury, venus, earth, mars, jupiter, saturn, uranus, neptune]

# 背景星空
stars = []
for _ in range(200):
    x = random.randint(0, WIDTH)
    y = random.randint(0, HEIGHT)
    size = random.randint(1, 2)
    brightness = random.randint(100, 255)
    stars.append((x, y, size, (brightness, brightness, brightness)))

# 信息显示
font = pygame.font.SysFont(None, 24)
info_text = [
    "太阳系模拟",
    "按 ↑ 加速，按 ↓ 减速",
    "按 R 重置模拟",
    "按 ESC 退出"
]

# 时间缩放因子
time_scale = 1.0

# 主循环
running = True
while running:
    for event in pygame.event.get():
        if event.type == QUIT:
            running = False
        elif event.type == KEYDOWN:
            if event.key == K_ESCAPE:
                running = False
            elif event.key == K_UP:
                time_scale *= 2  # 加速
            elif event.key == K_DOWN:
                time_scale /= 2  # 减速
            elif event.key == K_r:
                # 重置模拟
                for planet in planets:
                    planet.angle = random.uniform(0, 2 * math.pi)
                    planet.orbit_points = []
    
    # 清屏
    screen.fill(BLACK)
    
    # 绘制星空背景
    for star in stars:
        x, y, size, color = star
        pygame.draw.circle(screen, color, (x, y), size)
    
    # 绘制太阳
    pygame.draw.circle(screen, YELLOW, (WIDTH // 2, HEIGHT // 2), sun.radius)
    
    # 绘制太阳光晕
    for i in range(3):
        radius = sun.radius + 5 + i * 3
        pygame.draw.circle(screen, (255, 255, 0, 100), (WIDTH // 2, HEIGHT // 2), radius, 1)
    
    # 更新和绘制行星
    for planet in planets:
        for _ in range(int(time_scale)):
            planet.update()
        planet.draw(screen)
    
    # 显示信息
    for i, text in enumerate(info_text):
        text_surface = font.render(text, True, WHITE)
        screen.blit(text_surface, (10, 10 + i * 25))
    
    # 显示时间缩放因子
    time_text = font.render(f"时间缩放: {time_scale:.2f}x", True, WHITE)
    screen.blit(time_text, (WIDTH - time_text.get_width() - 10, 10))
    
    # 显示帮助信息
    help_text = font.render("使用箭头键控制速度，R键重置", True, WHITE)
    screen.blit(help_text, (WIDTH // 2 - help_text.get_width() // 2, HEIGHT - 30))
    
    # 更新显示
    pygame.display.flip()
    
    # 控制帧率
    clock.tick(FPS)

# 退出游戏
pygame.quit()
sys.exit()